import torch
from tqdm import tqdm
from args_test import conf_load
from load_test import load_mod_tok,ta_lo,infer,save_output

def main():
    conf = conf_load()
    dev = conf.device
    type = torch.bfloat16 if conf.type == "bfloat16" else torch.float32
    torch.manual_seed(conf.sed)
    
    model,teken= load_mod_tok(
        conf.model,
        conf.adapter,
        type,
        dev)

    teta = ta_lo(conf.in_way)
    
    op = [infer(model, teken, [a], type, dev) for a in tqdm(teta)]

    save_output(conf.outfile,op)

if __name__ == "__main__":
    main()
